# pip install langchain-mcp-adapters
from major.models_manager import chat_model
from langchain_mcp_adapters.client import MultiServerMCPClient
from langgraph.prebuilt import create_react_agent
from langchain.tools import tool
import datetime
client = MultiServerMCPClient(
    {
        "demo": {
            "url": "http://127.0.0.1:8000/sse",
            "transport": "sse",
        }
    },
)
@tool
def current_time():
    """返回时间的工具"""
    print("正在获取时间...")
    return datetime.datetime.now()

async def stream_response():
    tools = await client.get_tools()
    tools.append(current_time)

    agent = create_react_agent(
        chat_model.get_model(),
        tools
    )
    response = await agent.ainvoke(
        {"messages": [{"role": "user", "content": "2*3等于几"}]}
    )
    print(response["messages"][-1].content)

if __name__ == "__main__":
    import asyncio
    asyncio.run(stream_response())